/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <gtest/gtest.h>

#include <IGLU/texture_loader/TextureLoaderFactory.h>

#include "../../util/Common.h"

#include <IGLU/texture_loader/ktx1/TextureLoaderFactory.h>
#include <IGLU/texture_loader/ktx2/TextureLoaderFactory.h>
#include <IGLU/texture_loader/stb_hdr/TextureLoaderFactory.h>
#include <IGLU/texture_loader/stb_jpeg/TextureLoaderFactory.h>
#include <IGLU/texture_loader/stb_png/TextureLoaderFactory.h>
#include <array>

namespace igl::tests {

class TextureLoaderFactoryTest : public ::testing::Test {
 private:
 public:
  TextureLoaderFactoryTest() = default;
  ~TextureLoaderFactoryTest() override = default;

  void SetUp() override {
    util::createDeviceAndQueue(iglDev_, cmdQueue_);
    ASSERT_TRUE(iglDev_ != nullptr);
    ASSERT_TRUE(cmdQueue_ != nullptr);
  }

  static std::vector<std::unique_ptr<iglu::textureloader::ITextureLoaderFactory>>
  createLoaderFactories() {
    std::vector<std::unique_ptr<iglu::textureloader::ITextureLoaderFactory>> factories;
    factories.reserve(5);
    factories.emplace_back(std::make_unique<iglu::textureloader::stb::hdr::TextureLoaderFactory>());
    factories.emplace_back(
        std::make_unique<iglu::textureloader::stb::jpeg::TextureLoaderFactory>());
    factories.emplace_back(std::make_unique<iglu::textureloader::stb::png::TextureLoaderFactory>());
    factories.emplace_back(std::make_unique<iglu::textureloader::ktx1::TextureLoaderFactory>());
    factories.emplace_back(std::make_unique<iglu::textureloader::ktx2::TextureLoaderFactory>());

    return factories;
  }

  // Member variables
 protected:
  void runLoadTest(const uint8_t* data, uint32_t size, const char* tag);
  iglu::textureloader::TextureLoaderFactory factory_{createLoaderFactories()};
  std::shared_ptr<IDevice> iglDev_;
  std::shared_ptr<ICommandQueue> cmdQueue_;
};

void TextureLoaderFactoryTest::runLoadTest(const uint8_t* data, uint32_t size, const char* tag) {
  ASSERT_NE(data, nullptr) << tag;

  Result result;

  auto loader = factory_.tryCreate(data, size, &result);
  ASSERT_NE(loader, nullptr) << tag;
  ASSERT_TRUE(result.isOk()) << tag << " " << result.message;

  auto texture = loader->create(*iglDev_, &result);
  ASSERT_NE(texture, nullptr) << tag;
  ASSERT_TRUE(result.isOk()) << tag << " " << result.message;

  loader->upload(*texture, &result);
  ASSERT_TRUE(result.isOk()) << tag << " " << result.message;

  if (loader->shouldGenerateMipmaps()) {
    texture->generateMipmap(*cmdQueue_);
  }
}

namespace {
constexpr const std::array<uint8_t, 67> kSinglePixelPng{
    {0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48,
     0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x00, 0x00, 0x00,
     0x00, 0x3A, 0x7E, 0x9B, 0x55, 0x00, 0x00, 0x00, 0x0A, 0x49, 0x44, 0x41, 0x54, 0x08,
     0xD7, 0x63, 0x60, 0x00, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE2, 0x21, 0xBC, 0x33, 0x00,
     0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82}};
}

TEST_F(TextureLoaderFactoryTest, loadPng) {
  runLoadTest(kSinglePixelPng.data(), kSinglePixelPng.size(), "PNG");
}

namespace {

constexpr const std::array<uint8_t, 160> kSinglePixelJpg{
    {0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10, 0x4A, 0x46, 0x49, 0x46, 0x00, 0x01, 0x01, 0x01, 0x00,
     0x48, 0x00, 0x48, 0x00, 0x00, 0xFF, 0xDB, 0x00, 0x43, 0x00, 0x03, 0x02, 0x02, 0x03, 0x02,
     0x02, 0x03, 0x03, 0x03, 0x03, 0x04, 0x03, 0x03, 0x04, 0x05, 0x08, 0x05, 0x05, 0x04, 0x04,
     0x05, 0x0A, 0x07, 0x07, 0x06, 0x08, 0x0C, 0x0A, 0x0C, 0x0C, 0x0B, 0x0A, 0x0B, 0x0B, 0x0D,
     0x0E, 0x12, 0x10, 0x0D, 0x0E, 0x11, 0x0E, 0x0B, 0x0B, 0x10, 0x16, 0x10, 0x11, 0x13, 0x14,
     0x15, 0x15, 0x15, 0x0C, 0x0F, 0x17, 0x18, 0x16, 0x14, 0x18, 0x12, 0x14, 0x15, 0x14, 0xFF,
     0xC0, 0x00, 0x0B, 0x08, 0x00, 0x01, 0x00, 0x01, 0x01, 0x01, 0x11, 0x00, 0xFF, 0xC4, 0x00,
     0x14, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
     0x00, 0x00, 0x00, 0x09, 0xFF, 0xC4, 0x00, 0x14, 0x10, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
     0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xDA, 0x00, 0x08,
     0x01, 0x01, 0x00, 0x00, 0x3F, 0x00, 0x54, 0xDF, 0xFF, 0xD9}};
}

TEST_F(TextureLoaderFactoryTest, loadJpg) {
  runLoadTest(kSinglePixelJpg.data(), kSinglePixelJpg.size(), "JPG");
}

namespace {
constexpr const std::array<uint8_t, 63> kSinglePixelHdr{
    {'#',  '?', 'R', 'A', 'D', 'I', 'A', 'N', 'C',  'E', '\n', 'S',  'O',  'F',  'T', 'W',
     'A',  'R', 'E', '=', 'G', 'E', 'G', 'L', '\n', 'F', 'O',  'R',  'M',  'A',  'T', '=',
     '3',  '2', '-', 'b', 'i', 't', '_', 'r', 'l',  'e', '_',  'r',  'g',  'b',  'e', '\n',
     '\n', '-', 'Y', ' ', '1', ' ', '+', 'X', ' ',  '1', '\n', 0x80, 0x80, 0x80, 0x81}};
}

TEST_F(TextureLoaderFactoryTest, loadHdr) {
  runLoadTest(kSinglePixelHdr.data(), kSinglePixelHdr.size(), "HDR");
}

namespace {
constexpr const std::array<uint8_t, 72> kSinglePixelKtx1{
    0xab, 0x4b, 0x54, 0x58, 0x20, 0x31, 0x31, 0xbb, 0x0d, 0x0a, 0x1a, 0x0a, 0x01, 0x02, 0x03,
    0x04, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x58, 0x80,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
}

TEST_F(TextureLoaderFactoryTest, loadKtx1) {
  runLoadTest(kSinglePixelKtx1.data(), kSinglePixelKtx1.size(), "KTX1");
}

namespace {
constexpr const std::array<uint8_t, 200> kSinglePixelKtx2{
    0xab, 0x4b, 0x54, 0x58, 0x20, 0x32, 0x30, 0xbb, 0x0d, 0x0a, 0x1a, 0x0a, 0x25, 0x00, 0x00, 0x00,
    0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x68, 0x00, 0x00, 0x00, 0x5c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xc4, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x5c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x02, 0x00, 0x58, 0x00, 0x01, 0x01, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xff, 0x00, 0x00, 0x00, 0x08, 0x00, 0x07, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xff, 0x00, 0x00, 0x00, 0x10, 0x00, 0x07, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xff, 0x00, 0x00, 0x00, 0x18, 0x00, 0x07, 0x1f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
}

TEST_F(TextureLoaderFactoryTest, loadKtx2) {
  runLoadTest(kSinglePixelKtx2.data(), kSinglePixelKtx2.size(), "KTX2");
}

} // namespace igl::tests
